import sys
import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
import joblib
import os
import logging
from datetime import datetime

# 获取当前脚本的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# 添加父目录到 sys.path
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)
from config import MODEL_INPUT, MODEL_OUTPUT, PREDICTION_OUTPUT, SCALER_OUTPUT

# -------------------------------
# 日志配置
# -------------------------------
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# -------------------------------
# 主函数：训练流量预测模型
# -------------------------------
def train_traffic_forecast_model():
    # 设置所有随机种子确保结果可重现
    np.random.seed(42)  # 设置numpy随机种子
    # Python的hash种子
    os.environ['PYTHONHASHSEED'] = '42'
    # LightGBM的随机种子已在params中设置
    
    logger.info("开始加载特征数据...")
    try:
        df = pd.read_csv(MODEL_INPUT)
        logger.info(f"数据加载成功，形状: {df.shape}")
    except Exception as e:
        logger.error(f"数据加载失败: {e}")
        return

    # 创建调试目录：用于保存各阶段中间结果
    debug_dir = os.path.join(os.path.dirname(MODEL_OUTPUT))
    os.makedirs(debug_dir, exist_ok=True)

    # -------------------------------
    # 1 特征选择：预测 in_mbps_measured
    # -------------------------------
    target_col = 'in_mbps_measured'
    meaningful_features = [
        # 时间
        'hour_of_day', 'day_of_week', 'is_peak_hour', 'is_weekend',
        # 空间
        'is_core_region', 'adjacent_count',
        # 事件
        'ddos_attack_flag', 'fault_impact_score', 'maintenance_flag',
        # 流量状态
        'in_mbps_lag_1', 'out_mbps_lag_1',
        'in_mbps_rolling_mean_1h', 'in_mbps_rolling_std_1h',
        # 质量
        'device_in_mbps_zscore', 'in_mbps_error_ratio',
        # 协议（已编码）
        'proto_TCP', 'proto_UDP'
    ]

    # 确保这些特征都在数据中
    feature_cols = [col for col in meaningful_features if col in df.columns]
    
    X = df[feature_cols]
    y = df[target_col]

    # 过滤只保留数值类型列
    X = X.select_dtypes(include=['int64', 'float64'])

    # 检查是否为空
    if X.shape[1] == 0:
        raise ValueError("所有特征都被过滤！请检查 feature_cols 是否包含非数值列。")

    logger.info(f"最终用于建模的数值型特征数量: {X.shape[1]}")
    logger.info(f"使用 {len(feature_cols)} 个特征进行建模")
    logger.info(f"特征示例: {feature_cols[:5]}...")

    # 阶段1完成：保存 X 和 y
    stage1_data = pd.concat([X.reset_index(drop=True), y.reset_index(drop=True).rename('target')], axis=1)
    stage1_path = os.path.join(debug_dir, 'stage1_feature_selection.csv')
    stage1_data.to_csv(stage1_path, index=False)
    logger.info(f"阶段1完成 - 特征选择结果已保存: {stage1_path}")

    # -------------------------------
    # 2 数据集划分（时间顺序保留）
    # -------------------------------
    split_idx = int(len(X) * 0.8)
    X_train, X_test = X[:split_idx], X[split_idx:]
    y_train, y_test = y[:split_idx], y[split_idx:]

    logger.info(f"训练集: {X_train.shape}, 测试集: {X_test.shape}")

    # 阶段2完成：保存划分后的训练/测试集
    X_train.to_csv(os.path.join(debug_dir, 'stage2_X_train.csv'), index=False)
    X_test.to_csv(os.path.join(debug_dir, 'stage2_X_test.csv'), index=False)
    pd.DataFrame(y_train).to_csv(os.path.join(debug_dir, 'stage2_y_train.csv'), index=False)
    pd.DataFrame(y_test).to_csv(os.path.join(debug_dir, 'stage2_y_test.csv'), index=False)
    logger.info(f"阶段2完成 - 数据划分结果已保存")

    # -------------------------------
    # 3 特征标准化（LightGBM 不必须，但为后续模型兼容）
    # -------------------------------
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 保存 scaler
    joblib.dump(scaler, SCALER_OUTPUT)
    logger.info(f"标准化器已保存: {SCALER_OUTPUT}")

    # 阶段3完成：保存标准化后的数据（转为 DataFrame 更可读）
    pd.DataFrame(X_train_scaled, columns=X_train.columns).to_csv(os.path.join(debug_dir, 'stage3_X_train_scaled.csv'), index=False)
    pd.DataFrame(X_test_scaled, columns=X_test.columns).to_csv(os.path.join(debug_dir, 'stage3_X_test_scaled.csv'), index=False)
    logger.info(f"阶段3完成 - 标准化结果已保存")

    # -------------------------------
    # 4 LightGBM 模型训练
    # -------------------------------
    logger.info("开始训练 LightGBM 模型...")
    lgb_train = lgb.Dataset(X_train_scaled, label=y_train)
    lgb_eval = lgb.Dataset(X_test_scaled, label=y_test, reference=lgb_train)

    params = {
        'objective': 'regression',
        'metric': 'mae',
        'boosting_type': 'gbdt',
        'num_leaves': 31,
        'learning_rate': 0.05,
        'feature_fraction': 0.9,
        'bagging_fraction': 0.8,
        'bagging_freq': 5,
        'verbose': -1,
        'random_state': 42,  # 确保模型可重现
        'deterministic': True,  # 确保确定性行为
    }

    model = lgb.train(
        params,
        lgb_train,
        num_boost_round=1000,
        valid_sets=[lgb_train, lgb_eval],
        callbacks=[lgb.early_stopping(50), lgb.log_evaluation(100)]
    )

    # 保存模型 - 使用LightGBM原生格式而不是joblib
    model.save_model(MODEL_OUTPUT)  # 修改点：使用原生保存方法
    logger.info(f" 模型训练完成，已保存: {MODEL_OUTPUT}")

    # -------------------------------
    # 模型评估
    # -------------------------------
    y_pred = model.predict(X_test_scaled, num_iteration=model.best_iteration)
    mae = mean_absolute_error(y_test, y_pred)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    r2 = r2_score(y_test, y_pred)

    logger.info(f"模型评估结果:")
    logger.info(f"MAE:  {mae:.2f} Mbps")
    logger.info(f"RMSE: {rmse:.2f} Mbps")
    logger.info(f"R²:   {r2:.4f}")

    # -------------------------------
    # 5 保存预测结果
    # -------------------------------
    results_df = pd.DataFrame({
        'timestamp': df['timestamp'][split_idx:].values,
        'region_id': df['region_id'][split_idx:].values,
        'true_in_mbps': y_test.values,
        'pred_in_mbps': y_pred
    })
    results_df.to_csv(PREDICTION_OUTPUT, index=False)
    logger.info(f"预测结果已保存: {PREDICTION_OUTPUT}")

if __name__ == "__main__":
    train_traffic_forecast_model()